import glob
import json

import numpy as np
import os
import re
from os import listdir
from os.path import join
from scipy import io
import pandas as pd
from torch.utils.data import DataLoader, Dataset
from .process import *
import torch
from datasets.DVS_Lip.native.cvtransforms import *

repr_map = {'eventFrame': get_eventFrame,
            'eventAccuFrame': get_eventAccuFrame,
            'timeSurface': get_timeSurface,
            'eventCount': get_eventCount}


# left or right move all event locations randomly
def random_shift_events(events, max_shift=20, resolution=(180, 240)):
    H, W = resolution
    x_shift, y_shift = np.random.randint(-max_shift, max_shift + 1, size=(2,))
    events[:, 1] += x_shift
    events[:, 2] += y_shift

    valid_events = (events[:, 1] >= 0) & (events[:, 1] < W) & (events[:, 2] >= 0) & (events[:, 2] < H)
    events = events[valid_events]

    return events


# flip half of the event images along the x dimension
def random_flip_events_along_x(events, resolution=(180, 240), p=0.5):
    """
    event locations: (t, x, y, p)
    """
    H, W = resolution
    if np.random.random() < p:
        events[:, 1] = W - 1 - events[:, 1]
    return events


def events_to_voxel_grid_pytorch(events, num_bins, width, height, device):
    """
    Build a voxel grid with bilinear interpolation in the time domain from a set of events.
    :param events: a [N x 4] NumPy array containing one event per row in the form: [timestamp, x, y, polarity]
    :param num_bins: number of bins in the temporal axis of the voxel grid
    :param width, height: dimensions of the voxel grid
    :param device: device to use to perform computations
    :return voxel_grid: PyTorch event tensor (on the device specified)
    """

    assert (events.shape[1] == 4)
    assert (num_bins > 0)
    assert (width > 0)
    assert (height > 0)

    with torch.no_grad():

        events_torch = torch.from_numpy(events).float()
        events_torch = events_torch.to(device)

        voxel_grid = torch.zeros(num_bins, height, width, dtype=torch.float32, device=device)
        if events_torch.shape[0] == 0:
            return voxel_grid

        voxel_grid = voxel_grid.flatten()

        # normalize the event timestamps so that they lie between 0 and num_bins
        last_stamp = events_torch[-1, 0]
        first_stamp = events_torch[0, 0]
        deltaT = float(last_stamp - first_stamp)

        if deltaT == 0:
            deltaT = 1.0

        events_torch[:, 0] = (num_bins - 1) * (events_torch[:, 0] - first_stamp) / deltaT
        ts = events_torch[:, 0]
        xs = events_torch[:, 1].long()
        ys = events_torch[:, 2].long()
        pols = events_torch[:, 3].float()
        pols[pols == 0] = -1  # polarity should be +1 / -1

        tis = torch.floor(ts)
        tis_long = tis.long()
        dts = ts - tis
        vals_left = pols * (1.0 - dts.float())
        vals_right = pols * dts.float()

        valid_indices = tis < num_bins
        valid_indices &= tis >= 0
        voxel_grid.index_add_(dim=0,
                              index=xs[valid_indices] + ys[valid_indices]
                                    * width + tis_long[valid_indices] * width * height,
                              source=vals_left[valid_indices])

        valid_indices = (tis + 1) < num_bins
        valid_indices &= tis >= 0

        voxel_grid.index_add_(dim=0,
                              index=xs[valid_indices] + ys[valid_indices] * width
                                    + (tis_long[valid_indices] + 1) * width * height,
                              source=vals_right[valid_indices])

        voxel_grid = voxel_grid.view(num_bins, height, width)

    return voxel_grid


def events_to_voxel_all(events, frame_nums, seq_len, num_bins, width, height, device):
    voxel_len = min(seq_len, frame_nums) * num_bins
    voxel_grid_all = np.zeros((num_bins * seq_len, 1, height, width))
    voxel_grid = events_to_voxel_grid_pytorch(events, voxel_len, width, height, device)
    voxel_grid = voxel_grid.unsqueeze(1).cpu().numpy()
    voxel_grid_all[:voxel_len] = voxel_grid
    return voxel_grid_all


class DVS_Lip(Dataset):
    def __init__(self, datafile="../DVS_Lip", phase='train', augmentation=False,
                 repr=['timeSurface', 'eventFrame'], time_num=120,
                 frame_transform=None, label_transform=None, native=False):
        self.augmentation = augmentation
        self.datafile = datafile
        self.phase = phase
        self.repr = repr
        self.time_num = time_num
        self.native = native
        self.length = 30
        self.frame_transform = frame_transform
        self.label_transform = label_transform
        self.labels = sorted(os.listdir(os.path.join(datafile, phase)))

        self.file_list = sorted(glob.glob(os.path.join(self.datafile, phase, '*', '*.npy')))

        with open('./data/frame_nums.json', 'r') as f:
            self.frame_nums = json.load(f)

    def __len__(self):
        return len(self.file_list)

    def __getitem__(self, idx):
        """
        returns events and label, loading events from aedat
        :param idx:
        :return: x,y,t,p,  label
        events: (x,y,t,p)  x:[0,346] y:[0,260] p:0/1
        """
        # load timestamps
        word = self.file_list[idx].split('\\')[-2]
        person = self.file_list[idx].split('\\')[-1][:-4]
        frame_num = self.frame_nums[self.phase][word][int(person)]

        # load events
        try:
            events_input = np.load(self.file_list[idx])
        except:
            print(self.file_list[idx])
        events_input = events_input[np.where(
            (events_input['x'] >= 8) & (events_input['x'] < 120) & (events_input['y'] >= 8) & (
                    events_input['y'] < 120))]
        events_input['x'] -= 8
        events_input['y'] -= 8
        t, x, y, p = events_input['t'], events_input['x'], events_input['y'], events_input['p']
        events_input = np.stack([t, x, y, p], axis=-1)  # [N,4]
        label = self.labels.index(word)

        if self.native:
            # convert events to voxel_grid
            event_voxel_low = events_to_voxel_all(events_input, frame_num, self.length, self.time_num // self.length,
                                                  112, 112,
                                                  device='cpu')  # # (30*num_bins[1],1, 96, 96)
            event_voxel_high = events_to_voxel_all(events_input, frame_num, self.length, self.time_num // self.length,
                                                   112, 112,
                                                   device='cpu')  # # (30*num_bins[1],1, 96, 96)

            if self.phase == 'train':
                event_voxel_low, event_voxel_high = RandomCrop(event_voxel_low, event_voxel_high, (96, 96))
                event_voxel_low, event_voxel_high = HorizontalFlip(event_voxel_low, event_voxel_high)
            else:
                event_voxel_low, event_voxel_high = CenterCrop(event_voxel_low, event_voxel_high, (96, 96))

            reprs = np.concatenate((event_voxel_low, event_voxel_high), axis=1)  # (T,C, 96, 96)

        else:
            # normalize the timestamps
            _min = events_input[:, 0].min()
            _max = events_input[:, 0].max()
            events_input[:, 0] = (events_input[:, 0] - _min) / (_max - _min)

            if self.augmentation:
                events_input = random_shift_events(events_input)
                events_input = random_flip_events_along_x(events_input)

            reprs = []
            for repr_name in self.repr:
                # 输入(t,x,y,p), 输出(T,y,x) y:[0,260] x:[0:346]
                repr_array = repr_map[repr_name](events_input[:, 0], events_input[:, 1].astype(np.int32),
                                                 events_input[:, 2].astype(np.int32),
                                                 events_input[:, 3],
                                                 repr_size=(96, 96), time_num=self.time_num)

                # standardization
                mu = np.mean(repr_array)
                sigma = np.std(repr_array)
                repr_array = (repr_array - mu) / sigma

                reprs.append(repr_array)

            reprs = np.array(reprs)  # (C,T,260,346)
            reprs = np.transpose(reprs, (1, 0, 2, 3))  # (T,C,260,346)

        if self.frame_transform is not None:
            reprs = self.frame_transform(reprs)

        if self.label_transform is not None:
            label = self.label_transform(label)

        return reprs, label  # (T,C,260,346) (1)
